"""
Case Type   : DATAVEC模块
Case Name   : HNSW 索引修改并行度，并行构建、重建索引，功能正常
Create At   : 2025.4.2
Owner       : wang-tianjing1
Description :
    step1: 创建表并插入数据
    step2: 设置不同的并发值后create index
    step3: 设置不同的并发值后reindex
    step4: 清理环境
Expect      :
    step1: 成功
    step2: 成功，查询结果一致
    step3: 成功，查询结果一致
    step4: 成功
History     :
"""

import os
import unittest
from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node


class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        self.log = Logger()
        self.log.info(f'----{os.path.basename(__file__)}:start----')
        self.sh = CommonSH('PrimaryDbUser')
        self.pri_dbuser = Node(node='PrimaryDbUser')
        self.constant = Constant()
        self.com = Common()
        self.comsql = CommonSQL('PrimaryDbUser')
        self.tb_name = 't_datavec_index0118'
        self.fun_insert = 'f_datavec_index0118'
        self.param_settings = "set enable_seqscan=off; set enable_indexscan=on; set enable_bitmapscan =off;"

    def test_main(self):
        step_txt = '----step1: 创建表并插入数据; expect:成功----'
        self.log.info(step_txt)
        create_tb = f"drop table if exists {self.tb_name};" \
                    f"create table {self.tb_name} (id bigserial primary key, val vector(320));"
        result = self.sh.execut_db_sql(create_tb)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result, '执行失败:' + step_txt)

        step_txt = '----创建自定义插入函数; expect: 成功----'
        self.log.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        self.assertEqual(result, self.fun_insert, '执行失败: 创建自定义插入函数失败')

        step_txt = '----使用自定义函数插入数据; expect: 成功----'
        self.log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(320, 5000, '{self.tb_name}', 'val', 'vector');"
        result = self.sh.execut_db_sql(self.param_settings + insert_sql)
        self.assertIn('SUCCESS', result, '执行失败: 使用自定义函数插入数据失败')

        step_txt = '----step2: 设置不同的并发值后create index; expect: 查询结果一致----'
        self.log.info(step_txt)

        sql_1 = f"{self.param_settings}alter table {self.tb_name} set(parallel_workers=6);" \
                f"create index idx_{self.tb_name}_1 on {self.tb_name} using hnsw(val vector_l2_ops);" \
                f"select id from {self.tb_name} order by val<->array_fill(49,array[320])::vector limit 10;"
        result_1 = self.sh.execut_db_sql(sql_1)
        lines_1 = result_1.split('\n')
        query_result_1 = []
        in_result_1 = False
        for line in lines_1:
            if line.strip() == '':
                if in_result_1:
                    break
                continue
            if 'id' in line:
                in_result_1 = True
                continue
            if in_result_1:
                query_result_1.append(line.strip())
        final_result_1 = '\n'.join(query_result_1)

        sql_2 = f"{self.param_settings}alter table {self.tb_name} set(parallel_workers=12);" \
                f"drop index idx_{self.tb_name}_1;" \
                f"create index idx_{self.tb_name}_2 on {self.tb_name} using hnsw(val vector_l2_ops);" \
                f"select id from {self.tb_name} order by val<->array_fill(49,array[320])::vector limit 10;"
        result_2 = self.sh.execut_db_sql(sql_2)
        lines_2 = result_2.split('\n')
        query_result_2 = []
        in_result_2 = False
        for line in lines_2:
            if line.strip() == '':
                if in_result_2:
                    break
                continue
            if 'id' in line:
                in_result_2 = True
                continue
            if in_result_2:
                query_result_2.append(line.strip())
        final_result_2 = '\n'.join(query_result_2)

        sql_3 = f"{self.param_settings}alter table {self.tb_name} reset(parallel_workers);" \
                f"drop index idx_{self.tb_name}_2;" \
                f"create index idx_{self.tb_name}_3 on {self.tb_name} using hnsw(val vector_l2_ops);" \
                f"select id from {self.tb_name} order by val<->array_fill(49,array[320])::vector limit 10;"
        result_3 = self.sh.execut_db_sql(sql_3)
        lines_3 = result_3.split('\n')
        query_result_3 = []
        in_result_3 = False
        for line in lines_3:
            if line.strip() == '':
                if in_result_3:
                    break
                continue
            if 'id' in line:
                in_result_3 = True
                continue
            if in_result_3:
                query_result_3.append(line.strip())
        final_result_3 = '\n'.join(query_result_3)
        
        self.log.info(f"idx_{self.tb_name}_1: {final_result_1}, idx_{self.tb_name}_2: {final_result_2}, idx_{self.tb_name}_3: {final_result_3}")
        self.assertEqual(final_result_1, final_result_2, '执行失败: 不同并发值下create index查询结果不一致，第2次查询')
        self.assertEqual(final_result_1, final_result_3, '执行失败: 不同并发值下create index查询结果不一致，第3次查询')

        step_txt = '----step3: 设置不同的并发值后reindex; expect: 查询结果一致----'
        self.log.info(step_txt)

        sql_step3_1 = f"{self.param_settings}alter table {self.tb_name} set(parallel_workers=6);" \
                      f"create index idx_{self.tb_name}_1 on {self.tb_name} using hnsw(val vector_l2_ops);" \
                      f"alter index idx_{self.tb_name}_1 set (pq_m=32,pq_ksub=20);" \
                      f"reindex index idx_{self.tb_name}_1;" \
                      f"select id from {self.tb_name} order by val<->array_fill(49,array[320])::vector limit 10;"
        result_step3_1 = self.sh.execut_db_sql(sql_step3_1)
        lines_step3_1 = result_step3_1.split('\n')
        query_result_step3_1 = []
        in_result_step3_1 = False
        for line in lines_step3_1:
            if line.strip() == '':
                if in_result_step3_1:
                    break
                continue
            if 'id' in line:
                in_result_step3_1 = True
                continue
            if in_result_step3_1:
                query_result_step3_1.append(line.strip())
        final_result_step3_1 = '\n'.join(query_result_step3_1)

        sql_step3_2 = f"{self.param_settings}alter table {self.tb_name} set(parallel_workers=12);" \
                      f"drop index idx_{self.tb_name}_1;" \
                      f"create index idx_{self.tb_name}_1 on {self.tb_name} using hnsw(val vector_l2_ops);" \
                      f"alter index idx_{self.tb_name}_1 set (pq_m=32,pq_ksub=20);" \
                      f"reindex index idx_{self.tb_name}_1;" \
                      f"select id from {self.tb_name} order by val<->array_fill(49,array[320])::vector limit 10;"
        result_step3_2 = self.sh.execut_db_sql(sql_step3_2)
        lines_step3_2 = result_step3_2.split('\n')
        query_result_step3_2 = []
        in_result_step3_2 = False
        for line in lines_step3_2:
            if line.strip() == '':
                if in_result_step3_2:
                    break
                continue
            if 'id' in line:
                in_result_step3_2 = True
                continue
            if in_result_step3_2:
                query_result_step3_2.append(line.strip())
        final_result_step3_2 = '\n'.join(query_result_step3_2)

        sql_step3_3 = f"{self.param_settings}alter table {self.tb_name} reset(parallel_workers);" \
                      f"drop index idx_{self.tb_name}_1;" \
                      f"create index idx_{self.tb_name}_1 on {self.tb_name} using hnsw(val vector_l2_ops);" \
                      f"alter index idx_{self.tb_name}_1 set (pq_m=32,pq_ksub=20);" \
                      f"reindex index idx_{self.tb_name}_1;" \
                      f"select id from {self.tb_name} order by val<->array_fill(49,array[320])::vector limit 10;"
        result_step3_3 = self.sh.execut_db_sql(sql_step3_3)
        lines_step3_3 = result_step3_3.split('\n')
        query_result_step3_3 = []
        in_result_step3_3 = False
        for line in lines_step3_3:
            if line.strip() == '':
                if in_result_step3_3:
                    break
                continue
            if 'id' in line:
                in_result_step3_3 = True
                continue
            if in_result_step3_3:
                query_result_step3_3.append(line.strip())
        final_result_step3_3 = '\n'.join(query_result_step3_3)
        self.log.info(f"idx_{self.tb_name}_1: {final_result_step3_1}, idx_{self.tb_name}_2: {final_result_step3_2}, idx_{self.tb_name}_3: {final_result_step3_3}")
        self.assertEqual(final_result_step3_1, final_result_step3_2, '执行失败: 不同并发值下reindex查询结果不一致，第2次查询')
        self.assertEqual(final_result_step3_1, final_result_step3_3, '执行失败: 不同并发值下reindex查询结果不一致，第3次查询')

    def tearDown(self):
        step_txt = '----step4: 清理环境; expect:成功----'
        self.log.info(step_txt)
        drop_cmd = f"drop table if exists {self.tb_name};"
        result = self.sh.execut_db_sql(self.param_settings + drop_cmd)
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, result, '执行失败:' + step_txt)
        self.log.info(f'----{os.path.basename(__file__)}:end----')
